

import pandas as pd
from numpy import log, NAN
from sklearn import preprocessing


class FractalDimension(object):
    def __init__(self, data, num_days):
        self.data = data
        self.num_days = num_days

    def log_return(self, n):
        return_rate = log(self.data[n:].values / self.data[:-n].values)
        return_rate = pd.Series(return_rate, index=self.data.index[n:])
        return return_rate

    def forward(self):
        r = abs(self.log_return(n=1))
        r_n = abs(self.log_return(n=self.num_days))
        r_scaled = r.rolling(window=self.num_days).sum().iloc[self.num_days-1:].values / (r_n.values/self.num_days)
        fractal_dim = log(r_scaled) / log(self.num_days)
        fractal_dim = pd.Series(fractal_dim, index=r_n.index)
        return fractal_dim


def calculate_fractal_dimension(data, num_days):
    """
    

    Parameters
    ----------
    data : TYPE
        DESCRIPTION.series
    num_days : TYPE
        DESCRIPTION.int

    Returns
    -------
    xs : TYPE
        DESCRIPTION.series

    """
    xs = data.fillna(method='ffill')
    xs_index = xs.index
    xs = xs.values.reshape(-1,1)
    scaler = preprocessing.MinMaxScaler()
    xs = scaler.fit_transform(xs).flatten()
    xs = pd.Series(xs + 0.00001, index=xs_index)
    

    fd = FractalDimension(xs, num_days)
    xs = fd.forward()
    
    xs = [NAN] * num_days + xs.to_list()
    xs = pd.Series(xs, index=xs_index)

    return xs


def calculate_fd_by_time(data, num_days, end_datetime):
    xs = calculate_fractal_dimension(data, num_days)
    if not isinstance(end_datetime, str):
        end_datetime = str(end_datetime)[:10]
    return xs[end_datetime]


if __name__ == "__main__":
    data = pd.read_excel('D:/春节版/input/y集合.xlsx', index_col=0).iloc[:,3]
    fd = calculate_fractal_dimension(data, 12)
    a = calculate_fd_by_time(data, 12, '2018-12-31')